{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-13T12:50:49.894627Z",
     "iopub.status.busy": "2021-07-13T12:50:49.894089Z",
     "iopub.status.idle": "2021-07-13T12:50:51.076632Z",
     "shell.execute_reply": "2021-07-13T12:50:51.076059Z",
     "shell.execute_reply.started": "2021-07-13T12:50:49.894561Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import os, sys, glob, argparse\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "import time, datetime\n",
    "import pdb, traceback\n",
    "\n",
    "import cv2\n",
    "from PIL import Image\n",
    "from sklearn.model_selection import train_test_split, StratifiedKFold, KFold\n",
    "\n",
    "import torch\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data.dataset import Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-13T13:06:20.623722Z",
     "iopub.status.busy": "2021-07-13T13:06:20.623129Z",
     "iopub.status.idle": "2021-07-13T13:06:20.697903Z",
     "shell.execute_reply": "2021-07-13T13:06:20.696737Z",
     "shell.execute_reply.started": "2021-07-13T13:06:20.623673Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_jpg = glob.glob('./Datawhale_人脸情绪识别_数据集/train/*/*')\n",
    "np.random.shuffle(train_jpg)\n",
    "train_jpg = np.array(train_jpg)\n",
    "    \n",
    "class QRDataset(Dataset):\n",
    "    def __init__(self, img_path, transform=None):\n",
    "        self.img_path = img_path\n",
    "        if transform is not None:\n",
    "            self.transform = transform\n",
    "        else:\n",
    "            self.transform = None\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        start_time = time.time()\n",
    "        img = Image.open(self.img_path[index]).convert('RGB')\n",
    "        \n",
    "        lbl_dict = {'angry': 0,\n",
    "             'disgusted': 1,\n",
    "             'fearful': 2,\n",
    "             'happy': 3,\n",
    "             'neutral': 4,\n",
    "             'sad': 5,\n",
    "             'surprised': 6}\n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img)\n",
    "        \n",
    "        if 'test' in self.img_path[index]:\n",
    "            return img, torch.from_numpy(np.array(0))\n",
    "        else:\n",
    "            lbl_int = lbl_dict[self.img_path[index].split('/')[-2]]\n",
    "            return img, torch.from_numpy(np.array(lbl_int))\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-13T12:59:07.046836Z",
     "iopub.status.busy": "2021-07-13T12:59:07.046265Z",
     "iopub.status.idle": "2021-07-13T12:59:07.052800Z",
     "shell.execute_reply": "2021-07-13T12:59:07.052072Z",
     "shell.execute_reply.started": "2021-07-13T12:59:07.046790Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class XunFeiNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(XunFeiNet, self).__init__()\n",
    "                \n",
    "        model = models.resnet18(True)\n",
    "        model.avgpool = nn.AdaptiveAvgPool2d(1)\n",
    "        model.fc = nn.Linear(512, 7)\n",
    "        self.resnet = model\n",
    "         \n",
    "    def forward(self, img):        \n",
    "        out = self.resnet(img)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-13T13:06:47.941529Z",
     "iopub.status.busy": "2021-07-13T13:06:47.940955Z",
     "iopub.status.idle": "2021-07-13T13:06:47.955653Z",
     "shell.execute_reply": "2021-07-13T13:06:47.954789Z",
     "shell.execute_reply.started": "2021-07-13T13:06:47.941484Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def validate(val_loader, model, criterion):\n",
    "    model.eval()\n",
    "    acc1 = []\n",
    "    with torch.no_grad():\n",
    "        end = time.time()\n",
    "        for i, (input, target) in enumerate(val_loader):\n",
    "            input = input.cuda()\n",
    "            target = target.cuda()\n",
    "            \n",
    "            output = model(input)\n",
    "            loss = criterion(output, target)\n",
    "            acc1.append((output.argmax(1) == target).float().mean().item())\n",
    "\n",
    "        print(' * Val Acc@1 {0}'.format(np.mean(acc1)))\n",
    "        return np.mean(acc1)\n",
    "\n",
    "def predict(test_loader, model, tta=10):\n",
    "    model.eval()\n",
    "    \n",
    "    test_pred_tta = None\n",
    "    for _ in range(tta):\n",
    "        test_pred = []\n",
    "        with torch.no_grad():\n",
    "            end = time.time()\n",
    "            for i, (input, target) in enumerate(test_loader):\n",
    "                input = input.cuda()\n",
    "                target = target.cuda()\n",
    "\n",
    "                output = model(input)\n",
    "                output = output.data.cpu().numpy()\n",
    "\n",
    "                test_pred.append(output)\n",
    "        test_pred = np.vstack(test_pred)\n",
    "    \n",
    "        if test_pred_tta is None:\n",
    "            test_pred_tta = test_pred\n",
    "        else:\n",
    "            test_pred_tta += test_pred\n",
    "    \n",
    "    return test_pred_tta\n",
    "\n",
    "def train(train_loader, model, criterion, optimizer, epoch):\n",
    "    model.train()\n",
    "\n",
    "    end = time.time()\n",
    "    acc1 = []\n",
    "    for i, (input, target) in enumerate(train_loader):\n",
    "        input = input.cuda(non_blocking=True)\n",
    "        target = target.cuda(non_blocking=True)\n",
    "        output = model(input)\n",
    "        loss = criterion(output, target)\n",
    "\n",
    "        acc1.append((output.argmax(1) == target).float().mean().item())\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if i % 100 == 0:\n",
    "            print('Train: {0}'.format(np.mean(acc1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-13T13:01:29.634784Z",
     "iopub.status.busy": "2021-07-13T13:01:29.634402Z",
     "iopub.status.idle": "2021-07-13T13:03:38.865262Z",
     "shell.execute_reply": "2021-07-13T13:03:38.863820Z",
     "shell.execute_reply.started": "2021-07-13T13:01:29.634752Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch:  0\n",
      "Train: 0.07999999821186066\n",
      "Train: 0.22099009348024237\n",
      "Train: 0.24338307917414614\n",
      "Train: 0.2624584656716185\n",
      "Train: 0.27855360942737123\n",
      " * Val Acc@1 0.34900000631809236\n",
      "\n",
      "Epoch:  1\n",
      "Train: 0.41999998688697815\n",
      "Train: 0.36831682035238433\n",
      "Train: 0.3699502394122271\n",
      "Train: 0.37614616974445675\n",
      "Train: 0.38334163566331314\n",
      " * Val Acc@1 0.4400000077486038\n",
      "\n",
      "Epoch:  2\n",
      "Train: 0.3400000035762787\n",
      "Train: 0.41999998924755816\n",
      "Train: 0.4163183962824333\n",
      "Train: 0.42019932412625943\n",
      "Train: 0.4243890156323773\n",
      " * Val Acc@1 0.48600000627338885\n",
      "\n",
      "Epoch:  3\n",
      "Train: 0.4399999976158142\n",
      "Train: 0.45108909506608946\n",
      "Train: 0.44935322059920774\n",
      "Train: 0.45202656481353154\n",
      "Train: 0.4526184410079756\n",
      " * Val Acc@1 0.510000008046627\n",
      "\n",
      "Epoch:  4\n",
      "Train: 0.3799999952316284\n",
      "Train: 0.47148513587394564\n",
      "Train: 0.47343282201396886\n",
      "Train: 0.47362124895138596\n",
      "Train: 0.4766084647312426\n",
      " * Val Acc@1 0.5310000072419644\n",
      "\n",
      "Epoch:  5\n",
      "Train: 0.3999999761581421\n",
      "Train: 0.498613845888931\n",
      "Train: 0.4941293383712199\n",
      "Train: 0.49508304176140466\n",
      "Train: 0.4924189380725423\n",
      " * Val Acc@1 0.5340000076591969\n",
      "\n",
      "Epoch:  6\n",
      "Train: 0.4399999976158142\n",
      "Train: 0.5073267165387031\n",
      "Train: 0.5070646615763802\n",
      "Train: 0.5072425105841056\n",
      "Train: 0.5103241755183499\n",
      " * Val Acc@1 0.5330000078678131\n",
      "\n",
      "Epoch:  7\n",
      "Train: 0.4399999976158142\n",
      "Train: 0.5160395889589102\n",
      "Train: 0.5123382949117404\n",
      "Train: 0.5136212487751464\n",
      "Train: 0.5171570934708279\n",
      " * Val Acc@1 0.5510000076144934\n",
      "\n",
      "Epoch:  8\n",
      "Train: 0.5799999833106995\n",
      "Train: 0.5263366191694052\n",
      "Train: 0.5302487420501993\n",
      "Train: 0.5262458330373035\n",
      "Train: 0.5265835261255726\n",
      " * Val Acc@1 0.5490000075101853\n",
      "\n",
      "Epoch:  9\n",
      "Train: 0.41999998688697815\n",
      "Train: 0.5277227576416318\n",
      "Train: 0.5382089401358989\n",
      "Train: 0.5424584567150801\n",
      "Train: 0.542044872729261\n",
      " * Val Acc@1 0.570000006556511\n"
     ]
    }
   ],
   "source": [
    "skf = KFold(n_splits=5, random_state=233, shuffle=True)\n",
    "for flod_idx, (train_idx, val_idx) in enumerate(skf.split(train_jpg, train_jpg)):\n",
    "    \n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        QRDataset(train_jpg[train_idx][:],\n",
    "                transforms.Compose([\n",
    "                            transforms.RandomAffine(10),\n",
    "                            transforms.ColorJitter(hue=.05, saturation=.05),\n",
    "                            transforms.RandomHorizontalFlip(),\n",
    "                            transforms.RandomVerticalFlip(),\n",
    "                            transforms.ToTensor(),\n",
    "                            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "            ])\n",
    "        ), batch_size=50, shuffle=True, num_workers=5, pin_memory=True\n",
    "    )\n",
    "    \n",
    "    val_loader = torch.utils.data.DataLoader(\n",
    "        QRDataset(train_jpg[val_idx][:1000],\n",
    "                transforms.Compose([\n",
    "#                             transforms.Resize((256, 256)),\n",
    "                            # transforms.Resize((124, 124)),\n",
    "                            # transforms.RandomCrop((88, 88)),\n",
    "                            transforms.ToTensor(),\n",
    "                            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "            ])\n",
    "        ), batch_size=10, shuffle=False, num_workers=5, pin_memory=True\n",
    "    )\n",
    "        \n",
    "    model = XunFeiNet().cuda()\n",
    "    criterion = nn.CrossEntropyLoss().cuda()\n",
    "    optimizer = torch.optim.SGD(model.parameters(), 0.003)\n",
    "    best_acc = 0.0\n",
    "    for epoch in range(10):\n",
    "        print('\\nEpoch: ', epoch)\n",
    "\n",
    "        train(train_loader, model, criterion, optimizer, epoch)\n",
    "        val_acc = validate(val_loader, model, criterion)\n",
    "        \n",
    "        if val_acc > best_acc:\n",
    "            best_acc = val_acc\n",
    "            torch.save(model.state_dict(), './resnet18_fold{0}.pt'.format(flod_idx))\n",
    "            \n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-13T13:06:50.346825Z",
     "iopub.status.busy": "2021-07-13T13:06:50.346233Z",
     "iopub.status.idle": "2021-07-13T13:07:51.466196Z",
     "shell.execute_reply": "2021-07-13T13:07:51.465443Z",
     "shell.execute_reply.started": "2021-07-13T13:06:50.346777Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_jpg = glob.glob('./Datawhale_人脸情绪识别_数据集/test/*')\n",
    "test_jpg = np.array(test_jpg)\n",
    "test_jpg.sort()\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "        QRDataset(test_jpg,\n",
    "                transforms.Compose([\n",
    "                            transforms.RandomHorizontalFlip(),\n",
    "                            transforms.RandomVerticalFlip(),\n",
    "                            transforms.ToTensor(),\n",
    "                            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "            ])\n",
    "        ), batch_size=50, shuffle=False, num_workers=5, pin_memory=True\n",
    ")\n",
    "        \n",
    "model = XunFeiNet().cuda()\n",
    "model.load_state_dict(torch.load('resnet18_fold0.pt'))\n",
    "test_pred = predict(test_loader, model, 5)\n",
    "\n",
    "\n",
    "# test_csv = pd.DataFrame()\n",
    "# test_csv['ID'] = list(range(0, 3082))\n",
    "# test_csv['Label'] = np.argmax(test_pred, 1)\n",
    "# test_csv['Label'] = test_csv['Label'].map({1:'pos', 0:'neg'})\n",
    "# test_csv.to_csv('tmp.csv', index=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-13T13:08:24.418232Z",
     "iopub.status.busy": "2021-07-13T13:08:24.417670Z",
     "iopub.status.idle": "2021-07-13T13:08:24.433246Z",
     "shell.execute_reply": "2021-07-13T13:08:24.432607Z",
     "shell.execute_reply.started": "2021-07-13T13:08:24.418186Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "cls_name = np.array(['angry', 'disgusted', 'fearful', 'happy','neutral', 'sad', 'surprised'])\n",
    "submit_df = pd.DataFrame({'name': test_jpg, 'label': cls_name[test_pred.argmax(1)]})\n",
    "submit_df['name'] = submit_df['name'].apply(lambda x: x.split('/')[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-13T13:08:25.736803Z",
     "iopub.status.busy": "2021-07-13T13:08:25.736238Z",
     "iopub.status.idle": "2021-07-13T13:08:25.758999Z",
     "shell.execute_reply": "2021-07-13T13:08:25.758024Z",
     "shell.execute_reply.started": "2021-07-13T13:08:25.736758Z"
    }
   },
   "outputs": [],
   "source": [
    "submit_df = submit_df.sort_values(by='name')\n",
    "submit_df.to_csv('pytorch_submit.csv', index=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
